import os
import seaborn as sns 
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from scipy import stats
from statsmodels.stats.multitest import multipletests


def run_multiple_comparisons(df, alpha=0.05):
    """安全的假设检验执行"""
    # 确保分组列存在
    if 'Group' not in df.columns:
        print("缺少Group列，无法进行组间比较")
        return pd.DataFrame()
        
    # 仅处理数值特征
    features = df.select_dtypes(include=np.number).columns
    results = []
    
    asd_mask = df['Group'] == 'ASD'
    td_mask = df['Group'] == 'TD'
    
    for feature in features:
        if feature == 'Group': continue
            
        asd_data = df.loc[asd_mask, feature]
        td_data = df.loc[td_mask, feature]
        
        # 忽略缺失或无效数据
        if asd_data.isna().all() or td_data.isna().all():
            continue
            
        # 组间差异检验
        try:
            u_stat, p_value = stats.mannwhitneyu(asd_data.dropna(), td_data.dropna())
        except ValueError:
            p_value = 1.0
        
        # 效应量计算
        n_asd, n_td = len(asd_data), len(td_data)
        pooled_std = np.sqrt(((n_asd-1)*asd_data.std()**2 + (n_td-1)*td_data.std()**2) / (n_asd+n_td-2))
        cohens_d = (asd_data.mean() - td_data.mean()) / pooled_std if pooled_std != 0 else 0
        
        results.append({
            'Feature': feature,
            'P_Value': p_value,
            'Cohens_d': cohens_d
        })
    
    if not results:
        return pd.DataFrame()
    
    results_df = pd.DataFrame(results)
    
    # FDR校正多重比较
    p_vals = results_df['P_Value'].values
    _, p_adjusted, _, _ = multipletests(p_vals, alpha=alpha, method='fdr_bh')
    
    results_df['FDR_Adjusted_P'] = p_adjusted
    results_df['Significant'] = p_adjusted < alpha
    
    return results_df.sort_values(by='FDR_Adjusted_P')

def plot_significant_features(results_df, save_path=''):
    """绘制显著特征的效应量图（带保护）"""
    if results_df.empty or 'Significant' not in results_df.columns:
        print("无显著特征可绘制")
        return
        
    sig_df = results_df[results_df['Significant']].sort_values('Cohens_d', ascending=False)
    
    if sig_df.empty:
        print("校正后无显著特征")
        return
    
    plt.figure(figsize=(10, 6))
    ax = sns.barplot(x='Cohens_d', y='Feature', data=sig_df, dodge=False)
    plt.axvline(0, color='k', linestyle='--')
    plt.xlabel("Cohen's d (效应量)")
    plt.ylabel('特征')
    plt.title(f'显著特征效应量 (FDR校正后P < 0.05)')
    
    # 添加效应量值标签
    for i in ax.containers:
        ax.bar_label(i, label_type='edge', fmt='%.2f', padding=5)
    
    plt.tight_layout()
    
    if save_path:
        output_file = os.path.join(save_path, 'significant_features.png')
        plt.savefig(output_file, dpi=300)
        print(f"保存效应量图: {output_file}")
    plt.close()